package com.example.demo.config;

import com.example.demo.Mapper.ChatMemoryMapper;
import com.example.demo.Mapper.MessageMapper;
import com.example.demo.Mapper.SessionMapper;
import com.example.demo.context.BaseContext;
import com.example.demo.service.*;
import com.example.demo.func.HighLevelCalculator;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;

import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.tool.ToolProvider;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.web.search.WebSearchTool;
import dev.langchain4j.web.search.searchapi.SearchApiWebSearchEngine;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.context.annotation.RequestScope;


import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

@Configuration
@RequiredArgsConstructor
@Slf4j
public class AssistantConfig {
    final ChatLanguageModel chatLanguageModel;

    final StreamingChatLanguageModel streamingChatLanguageModel;

//    @Bean
//    public EmbeddingStore<TextSegment> initEmbeddingStore(){
//        return new InMemoryEmbeddingStore<>();
//    }

    final MessageMapper messageMapper;

    final SessionMapper sessionMapper;

    final ChatMemoryMapper chatMemoryMapper;
    public PostgresChatMemoryStore postgresChatMemoryStore(){
        return new PostgresChatMemoryStore(chatMemoryMapper);
    }

    final EmbeddingModel embeddingModel;


    //SearchApiWebSearchEngine engine
    @Bean
    //@RequestScope
    public Assistant assistant(EmbeddingStore<TextSegment> embeddingStore){
        return AiServices.builder(Assistant.class)
                //配置上下文机制，保持最近10条的对话历史
                .chatMemoryProvider( memoryId-> MessageWindowChatMemory.withMaxMessages(10))
                //rag检索模块
                .contentRetriever(EmbeddingStoreContentRetriever.builder()
                        //指定向量存储的位置
                        .embeddingStore(embeddingStore)
                        //指定进行query使用的模型，进一步将用户输入转换为embedding向量，与文档向量进行相似的匹配
                        .embeddingModel(embeddingModel)
                        //.filter(metadataKey("userId").isEqualTo(BaseContext.getCurrentId()))
                        .build())
                //注入工具
                //.tools(new HighLevelCalculator(),new WebSearchTool(engine))
                .tools(new HighLevelCalculator())
                //配置chatModel，进行最终的回答
                .chatLanguageModel(chatLanguageModel).build();
    }

    @Bean
    public CommonAssistant commonAssistant() {
        return AiServices.builder(CommonAssistant.class)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.builder()
                        .id(memoryId)
                        .maxMessages(10)
                        .chatMemoryStore(postgresChatMemoryStore()) // 这里用 Spring 管理的 store
                        .build())
                .tools(new HighLevelCalculator())
                .chatLanguageModel(chatLanguageModel).build();
    }

    /**
     * 总结对话助手
     * @return
     */
    @Bean
    public SummaryAssistant summaryAssistant(){
        return AiServices.builder(SummaryAssistant.class)
                .chatLanguageModel(chatLanguageModel).build();
    }

    /**
     * 描述拓写助手
     * @return
     */
    @Bean
    public DescriptionAssistant descriptionAssistant(){
        return AiServices.builder(DescriptionAssistant.class)
                .chatLanguageModel(chatLanguageModel).build();
    }

    /**
     * 声色选择助手
     * @return
     */
    @Bean
    public VoiceAssistant voiceAssistant(){
        return AiServices.builder(VoiceAssistant.class)
                .chatLanguageModel(chatLanguageModel).build();
    }

    /**
     * 联网搜索助手
     * @param engine
     * @return
     */
    @Bean
    public WebAssistant webAssistant(SearchApiWebSearchEngine engine){
        return AiServices.builder(WebAssistant.class)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.builder()
                        .id(memoryId)
                        .maxMessages(10)
                        .chatMemoryStore(postgresChatMemoryStore()) // 这里用 Spring 管理的 store
                        .build())
                .tools(new WebSearchTool(engine))
                .streamingChatLanguageModel(streamingChatLanguageModel)
                .build();
    }


    @Bean
    @RequestScope
    public RagAssistant ragAssistant(EmbeddingStore<TextSegment> embeddingStore) {  // 这里注入 store
        return AiServices.builder(RagAssistant.class)
                //配置上下文机制，保持最近10条的对话历史
                .chatMemoryProvider( memoryId-> MessageWindowChatMemory.withMaxMessages(10))
                //rag检索模块
                .contentRetriever(EmbeddingStoreContentRetriever.builder()
                        //指定向量存储的位置
                        .embeddingStore(embeddingStore)
                        //指定进行query使用的模型，进一步将用户输入转换为embedding向量，与文档向量进行相似的匹配
                        .embeddingModel(embeddingModel)
                        //限制只获取用户id的文档
                        .filter(metadataKey("userId").isEqualTo(BaseContext.getCurrentId()))
                        .build())
                //注入工具
                //.tools(new HighLevelCalculator(),new WebSearchTool(engine))
                .tools(new HighLevelCalculator())
                //配置chatModel，进行最终的回答
                .streamingChatLanguageModel(streamingChatLanguageModel)
                .build();
    }

    /**
     * 流式回复的普通ai助手
     * @return
     */
    @Bean
    public StreamChatAssistant streamChatAssistant(SearchApiWebSearchEngine engine){
        return AiServices.builder(StreamChatAssistant.class)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.builder()
                        .id(memoryId)
                        .maxMessages(10)
                        .chatMemoryStore(postgresChatMemoryStore()) // 这里用 Spring 管理的 store
                        .build())
                .streamingChatLanguageModel(streamingChatLanguageModel)
                .tools(engine)
                .build();

    }
}
